import datetime

from yolo import YOLO
from PIL import Image
from utils.utils_gui import gen_name

image_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo']  # 可检测图片格式

data_path = "data/test.jpg"
save_path = "output/images/"
# -------------------------------------------------------------------------#
#   crop                指定了是否在单张图片预测后对目标进行截取
#   count               指定了是否进行目标的计数
# -------------------------------------------------------------------------#
crop = False
count = False

conf_thres = 0.01
iou_thres = 0.45

data_format = data_path.split(".")[-1]


# 检测图片
if data_format in image_formats:
    model = YOLO(cuda=True)
    image = Image.open(data_path)
    image = model.detect_image(image, conf_thres=conf_thres, iou_thres=iou_thres, crop=crop, count=count)
    image.show()
    image_output = save_path + gen_name(suffix=".jpg")
    image.save(image_output)
    print("检测结果已保存至" + image_output)
else:
    print("图片格式或者图片本身有问题，无法检测！！！")
